package tcpsession

import (
	"bytes"
	"encoding/binary"
	"errors"
	log "github.com/sirupsen/logrus"
	"go.uber.org/atomic"
	"io"
	"net"
	"time"
	"usbipproxy/common"
	"usbipproxy/domain"
	"usbipproxy/pb"
)

type TcpSession struct {
	domain.MsgPipeRepo
	connId int64
	c      net.Conn
	closed atomic.Bool
}

// readBinaryData 读取带长度前缀的二进制数据
func (t *TcpSession) readBinaryData() (*common.PacketHeader, []byte, error) {
	conn := t.c

	var header common.PacketHeader

	// 读取4字节的长度前缀
	lenBuf := make([]byte, common.HeaderSize)
	n, err := io.ReadFull(conn, lenBuf)
	if err != nil {
		log.Debugf("TcpSession readBinaryData err:%v", err)
		return nil, nil, err
	}
	if n != common.HeaderSize {
		log.Debugf("TcpSession readBinaryData n:%v not equal %v", n, common.HeaderSize)
	}

	if err := binary.Read(bytes.NewReader(lenBuf), binary.BigEndian, &header); err != nil {
		log.Warnf("TcpSession header pause error: %v", err)
		return nil, nil, err
	}

	// 验证魔数
	if header.Magic != common.MagicValue {
		log.Warnf("TcpSession invalid magic: 0x%X (should be 0x%X) header:%+v",
			header.Magic, common.MagicValue, header)
		return nil, nil, errors.New("invalid magic")
	}
	log.Debugf("TcpSession recv head: Cmd=%d, DataLen=%d", header.Command, header.DataLen)

	// 将长度前缀转换为uint32
	dataLen := header.DataLen

	if dataLen == 0 {
		return &header, nil, nil
	}

	// 限制最大数据长度（10MB）
	if dataLen > 10*1024*1024 {
		return nil, nil, errors.New("data too large")
	}

	// 读取实际数据
	data := make([]byte, dataLen)
	if _, err := io.ReadFull(conn, data); err != nil {
		return nil, nil, err
	}

	return &header, data, nil
}

func (t *TcpSession) displayInMsg(msg *pb.ProxyMsg) {
	if msg == nil {
		return
	}

	if msg.GetMsgType() != pb.MsgType_TO_CLIENT_ATTACH_DATA && msg.GetMsgType() != pb.MsgType_TO_SERVER_ATTACH_DATA {
		log.Infof("TcpSession recv <<==== msg:%v", msg)
	} else {
		if msg.GetMsgType() == pb.MsgType_TO_CLIENT_ATTACH_DATA {
			log.Debugf("TcpSession recv <<==== client_attach_data busId:%v sessId:%v len(data):%d ",
				msg.GetToClientAttachDataParam().GetBusId(), msg.GetToClientAttachDataParam().GetSessId(),
				len(msg.GetToClientAttachDataParam().GetData()))
		} else {
			log.Debugf("TcpSession recv <<==== server_attach_data busId:%v sessId:%v len(data):%d ",
				msg.GetToClientAttachDataParam().GetBusId(), msg.GetToClientAttachDataParam().GetSessId(),
				len(msg.GetToClientAttachDataParam().GetData()))
		}

	}
}

func (t *TcpSession) displayOutMsg(msg *pb.ProxyMsg) {
	if msg == nil {
		return
	}

	if msg.GetMsgType() != pb.MsgType_TO_CLIENT_ATTACH_DATA && msg.GetMsgType() != pb.MsgType_TO_SERVER_ATTACH_DATA {
		log.Infof("TcpSession send ====>> msg:%v", msg)
	} else {
		if msg.GetMsgType() == pb.MsgType_TO_CLIENT_ATTACH_DATA {
			log.Debugf("TcpSession send ====>> client_attach_data busId:%v sessId:%v len(data):%d ",
				msg.GetToClientAttachDataParam().GetBusId(), msg.GetToClientAttachDataParam().GetSessId(),
				len(msg.GetToClientAttachDataParam().GetData()))
		} else {
			log.Debugf("TcpSession send ====>> server_attach_data busId:%v sessId:%v len(data):%d ",
				msg.GetToClientAttachDataParam().GetBusId(), msg.GetToClientAttachDataParam().GetSessId(),
				len(msg.GetToClientAttachDataParam().GetData()))
		}

	}
}

// 处理接收到的包
func (t *TcpSession) handlePacket(command uint16, data []byte) {
	switch command {
	case common.CmdJson:
		log.Debugf("TcpSession handlePacket CmdJson recv len(data):%v", len(data))
		recvMsg, err := common.UnmarshalByType(common.TypeJson, data)
		if err != nil {
			log.Warnf("TcpSession handlePacket CmdJson parse msg failed. msg: %s err:%v", string(data), err)
			break
		}
		t.MsgPipeRepo.RecvInMsg(recvMsg)
		t.displayInMsg(recvMsg)
	case common.CmdData:
		log.Debugf("TcpSession handlePacket CmdData recv len(data):%v", len(data))
		recvMsg, _, err := common.Unmarshal(data)
		if err != nil {
			log.Warnf("TcpSession handlePacket parse msg failed. msg: %s err:%v", string(data), err)
			break
		}
		t.MsgPipeRepo.RecvInMsg(recvMsg)
		t.displayInMsg(recvMsg)
	case common.CmdHeartbeat:
		log.Debugf("TcpSession handlePacket heartbeat")

	default:
		log.Debugf("TcpSession handlePacket unkonwn cmd: %d", command)
	}
}

func (t *TcpSession) runRead() {
	log.Infof("TcpSession runRead goroutine start")

	for {
		header, buf, err := t.readBinaryData()
		if err != nil {
			log.Errorf("TcpSession read binary data err: %v", err)
			break
		} else {
			t.handlePacket(header.Command, buf)
		}
	}

	t.MsgPipeRepo.SendOutMsg(nil)
}

func (t *TcpSession) Start() error {
	defer func() {
		t.MsgPipeRepo.CloseIn()
		log.Println("connection read channel closed. ")
	}()
	go t.runWrite()
	t.runRead()
	return nil
}

func (t *TcpSession) runWrite() {
	log.Infof("TcpSession runWrite goroutine start")
	defer func() {
		t.c.Close()
		log.Infof("connection write channel closed. ")
	}()

	writeTimeout := 30 * time.Second
	for msg := range t.MsgPipeRepo.GetOutMsg() {
		if msg == nil {
			log.Infof("TcpSession runWrite msg is nil")
			break
		}

		data, err := common.Marshal(msg, common.TypePb)
		if err != nil {
			log.Warnf("TcpSession runWrite marshal msg failed. msg: %+v err:%v\n", msg, err)
			continue
		}

		//包装
		dataByte := common.PackData(common.CmdData, data)
		deadline := time.Now().Add(writeTimeout)
		if err := t.c.SetWriteDeadline(deadline); err != nil {
			log.Warnf("SetWriteDeadline fail: %v", err)
			break
		}
		if n, err := common.BlockingWrite(t.c, dataByte); err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				log.Warnf("TcpSession runWrite timeout %v", err)
			}
			log.Warnf("TcpSession runWrite write to client failed. msg: %v err:%v\n", msg, err)
			break
		} else {
			log.Debugf("TcpSession runWrite send length:%v data_len:%v", n, len(dataByte))
		}
		t.displayOutMsg(msg)
	}
	//_ = c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "OK"))
}

func (t *TcpSession) Stop() error {
	if !t.closed.CompareAndSwap(false, true) {
		return nil
	}
	t.MsgPipeRepo.SendOutMsg(nil)
	return nil
}

func (t *TcpSession) GetConnId() int64 {
	if t == nil {
		return 0
	}

	return t.connId
}

func NewTcpSession(connId int64, conn net.Conn, msgPipe domain.MsgPipeRepo) domain.TcpConnSessRepo {
	return &TcpSession{
		connId:      connId,
		c:           conn,
		MsgPipeRepo: msgPipe,
	}
}
